// Copyright 2022 <mzh.scnu@qq.com>. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"flag"
	"fmt"
	"strings"

	"google.golang.org/protobuf/compiler/protogen"
)

func main() {
	g := rpc{}

	var flags flag.FlagSet
	genCli := flags.Bool("gen-cli", false, "--tinyrpc_out=gen-cli=true:.")

	protogen.Options{
		ParamFunc: flags.Set,
	}.Run(func(p *protogen.Plugin) error {
		err := g.Generate(p)
		if err != nil {
			return err
		}
		if *genCli {
			return generateCli(p)
		}
		return nil
	})
}

type rpc struct{}

// Generate generate service code
func (md *rpc) Generate(plugin *protogen.Plugin) error {
	for _, f := range plugin.Files {
		if len(f.Services) == 0 {
			continue
		}
		fileName := f.GeneratedFilenamePrefix + ".svr.go"
		t := plugin.NewGeneratedFile(fileName, f.GoImportPath)
		t.P("// Code generated by protoc-gen-tinyrpc.")
		t.P()
		pkg := fmt.Sprintf("package %s", f.GoPackageName)
		t.P(pkg)
		t.P()
		for _, s := range f.Services {
			serviceCode := fmt.Sprintf(`%stype %s struct{}`,
				getComments(s.Comments), s.Desc.Name())
			t.P(serviceCode)
			t.P()
			for _, m := range s.Methods {
				funcCode := fmt.Sprintf(`%sfunc(this *%s) %s(args *%s,reply *%s)error{
					// define your service ...
					return nil
				}
				`, getComments(m.Comments), s.Desc.Name(),
					m.Desc.Name(), m.Input.Desc.Name(), m.Output.Desc.Name())
				t.P(funcCode)
			}
		}
	}
	return nil
}

func generateCli(plugin *protogen.Plugin) error {
	for _, f := range plugin.Files {
		if len(f.Services) == 0 {
			continue
		}
		fileName := f.GeneratedFilenamePrefix + ".cli.go"
		t := plugin.NewGeneratedFile(fileName, f.GoImportPath)
		t.P("// Code generated by protoc-gen-tinyrpc. DO NOT EDIT.")
		t.P()
		pkg := fmt.Sprintf("package %s", f.GoPackageName)
		t.P(pkg)
		t.P()
		importCode := `import (
		"github.com/zehuamama/tinyrpc"
		"net"
		)`
		t.P(importCode)
		t.P()
		for _, s := range f.Services {
			clientCode := fmt.Sprintf(`%stype %sClient struct{
				client *tinyrpc.Client
			}
			`,
				getComments(s.Comments), s.Desc.Name())
			t.P(clientCode)
			t.P()
			newClientCode := fmt.Sprintf(`func New%sClient(conn net.Conn) %sClient {
				return %sClient{client: tinyrpc.NewClient(conn)}
			}
			`,
				s.Desc.Name(), s.Desc.Name(), s.Desc.Name())
			t.P(newClientCode)
			t.P()
			for _, m := range s.Methods {
				funcCode := fmt.Sprintf(`%sfunc(this *%sClient) %s(args *%s,reply *%s)error{
					return this.client.Call("%s.%s", args, reply)
				}
				`, getComments(m.Comments), s.Desc.Name(),
					m.Desc.Name(), m.Input.Desc.Name(), m.Output.Desc.Name(), s.Desc.Name(), m.Desc.Name())
				t.P(funcCode)
			}
		}
	}
	return nil
}

// getComments get comment details
func getComments(comments protogen.CommentSet) string {
	c := make([]string, 0)
	c = append(c, strings.Split(string(comments.Leading), "\n")...)
	c = append(c, strings.Split(string(comments.Trailing), "\n")...)

	res := ""
	for _, comment := range c {
		if strings.TrimSpace(comment) == "" {
			continue
		}
		res += "//" + comment + "\n"
	}
	return res
}
